{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Foundational Model Fine-tuning using TAO Classification PyT\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is DinoV2 ?\n",
    "\n",
    "NV-Dinov2 is a visual foundational model trained on NVIDIA proprietary large scale dataset. DinoV2 is a self-supervised learning method that uses a combination of two SSL techniques : DINO and iBOT. These models could greatly simplify the use of images in any system by producing all purpose visual features, i.e., features that work across image distributions and tasks without finetuning. Trained on large curated datasets, our model has learnt robust fine-grained representation useful for localization and classification tasks. This model can be used as a foundation model for a variety of downstream tasks with few labeled examples.\n",
    "\n",
    "## What is CLIP ?\n",
    "\n",
    "CLIP (Contrastive Language-Image Pretraining) is a deep learning model developed by OpenAI. It's designed to understand images and text together in a way that allows it to perform a wide array of tasks. Unlike traditional computer vision models that are trained solely on images, or natural language models that are trained only on text, CLIP is trained on a large dataset containing both images and their associated textual descriptions.\n",
    "\n",
    "This allows CLIP to perform tasks like image classification, text-based image retrieval, and even generate textual descriptions for images. The key innovation is the use of a contrastive learning framework, which helps the model learn to associate images and their descriptions. Refer [here](https://openai.com/research/clip) for more information. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Finetune a vit_b_32 Open CLIP model on the ImageNet-1k dataset\n",
    "* Evaluate the trained model.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to a .onnx file for deployment to DeepStream.\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `classification` model\n",
    "which you may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of Classification using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Inferences](#head-6)\n",
    "7. [Deploy](#head-7)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=FIXME\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()))\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"cls_pyt_fm\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "# Point to the 'deps' folder in samples from where you are launching notebook inside classification folder.\n",
    "os.environ[\"PROJECT_DIR\"]=FIXME\n",
    "# Set your encryption key, and use the same key for all commands\n",
    "%env NUM_GPUS = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       },\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 525+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Prepare dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** This Notebook example is for 1000 classes of ImageNet. If you are using your custom dataset other than ImageNet - Please update the `dataset.data` config with `classes` field that points to a file with class names. Please refer to documentation for more details on the classes text file. Update the `num_classes` under `model.head` accordingly. For reference: Please refer to the `train_cats_dogs.yaml` in specs of `clsasification_pyt` under parent directory which gives an example of fine-tuning on 2-classes dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You need download the ImageNet2012 dataset and format it into train/ val/ test folders. The train, val folders should be unzipped and placed in $HOST_DATA_DIR/imagenet.\n",
    "\n",
    "The Data can be Downloaded by following instructions here: \n",
    "[MMPretrain Imagenet Download Instructions](https://mmpretrain.readthedocs.io/en/latest/user_guides/dataset_prepare.html) \n",
    "\n",
    "Go to official [Download page](http://www.image-net.org/download-images). Find download links for ILSVRC2012 and download the following two files.\n",
    "\n",
    "* ILSVRC2012_img_train.tar (~138GB)\n",
    "\n",
    "    * For training untar the class folders into the `train` such that the train folder has 1000 folders corresponding to each class.\n",
    "\n",
    "* ILSVRC2012_img_val.tar (~6.3GB)\n",
    "    * For validation images: You need to move the images to respective class folders. You can use this script for the same [valprep](https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh).\n",
    "\n",
    "You can also use this shell script to perform the above 2 steps: [extract_ILSVRC script](https://github.com/pytorch/examples/blob/main/imagenet/extract_ILSVRC.sh).\n",
    "Example contents of the final train/ val folder looks like this:\n",
    "\n",
    "**./train**\n",
    "\n",
    "    n07693725\n",
    "    ...\n",
    "    n07614500\n",
    "\n",
    "**./val**\n",
    "\n",
    "    n07693725\n",
    "    ...\n",
    "    n07614500\n",
    "\n",
    "\n",
    "The above steps can also be performed by the following bash commands, if you have not downloaded the ImageNet 2012 dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -P $HOST_DATA_DIR - https://raw.githubusercontent.com/pytorch/examples/main/imagenet/extract_ILSVRC.sh\n",
    "!wget -P $HOST_DATA_DIR - https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_train.tar\n",
    "!wget -P $HOST_DATA_DIR - https://image-net.org/data/ILSVRC/2012/ILSVRC2012_img_val.tar\n",
    "!(cd $HOST_DATA_DIR ; sh extract_ILSVRC.sh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the following dependencies for running the dataset preparation scripts\n",
    "!pip3 install Cython==0.29.36\n",
    "!pip3 install -r $PROJECT_DIR/requirements-pip.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Verify downloaded dataset <a class=\"anchor\" id=\"head-1-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $HOST_DATA_DIR/imagenet/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Please run this cell to ensure that your data has been downloaded properly\n",
    "import os \n",
    "train_dir = os.path.join(os.environ[\"LOCAL_PROJECT_DIR\"], \"imagenet/train\")\n",
    "if len(os.listdir(train_dir)) == 1000: \n",
    "    print(\"Successfully ImageNet Dataset Found.\") \n",
    "else: \n",
    "    print(\"Dataset Not Found. Please check properly.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* dataset:\n",
    "  * data:\n",
    "    * samples_per_gpu: Number of samples in a batch\n",
    "    * workers_per_gpu: workers per GPU\n",
    "    * train:\n",
    "      data_prefix: /data/imagenet/train\n",
    "      pipeline: Augmentations Config\n",
    "   * val:\n",
    "     * data_prefix: /data/imagenet/val\n",
    "   * test:\n",
    "     * data_prefix: /data/imagenet/val\n",
    "\n",
    "* model:\n",
    "  * backbone:\n",
    "    * type: \"open_clip\"\n",
    "    * custom_args:\n",
    "      * model_name: Model arch\n",
    "    * freeze: true\n",
    "    * pretrained: the pretrained dataset\n",
    "  * head:\n",
    "    * type: LinearClsHead\n",
    "    * num_classes: number of classes\n",
    "    * in_channels: 512\n",
    "    * loss: loss config\n",
    "\n",
    "* train:\n",
    "  * train_config:\n",
    "    * find_unused_parameters: True\n",
    "    * optimizer: Optimizer Config\n",
    "    * lr_config: Learning Rate Config \n",
    "    * optimizer_config: Optimizer Config\n",
    "    * runner:\n",
    "      * max_epochs: Max num of epochs to train\n",
    "    * checkpoint_config:\n",
    "      * interval: Intervals at which to save the checkpoint\n",
    "    * logging:\n",
    "      * interval: Intervals to do logging\n",
    "    * evaluation:\n",
    "      * interval: Interval at which to do evaluation\n",
    "\n",
    "Please refer to the TAO documentation about Classification to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/train_imagenet_clip.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete\n",
    "\n",
    "\n",
    "**Note:** If you are using NV-Dinov2 pre-trained weights from the NVAIE, the following parameters from the spec file model config should be used:\n",
    "\n",
    "* model:\n",
    "  * backbone:\n",
    "    * type: \"vit_large_patch14_dinov2_swiglu\"\n",
    "    * freeze: true\n",
    "    * pretrained: \"/path/to/NV_DINOV2_518.pth\"\n",
    "  * head:\n",
    "    * type: LinearClsHead\n",
    "    * num_classes: 1000\n",
    "    * in_channels: 1024\n",
    "    * loss:\n",
    "      * type: CrossEntropyLoss\n",
    "      * loss_weight: 1.0\n",
    "      * use_soft: False\n",
    "    * topk: [1, 5]\n",
    "\n",
    "\n",
    "**Note:** If you are using your custom dataset other than ImageNet - Please update the `dataset.data` config with `classes` field that points to a file with classes. Update the `num_classes` under `model.head` accordingly. For more details: Please refer to the `train_imagenet_clip.yaml` for reference. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List of Supported Backbones\n",
    "\n",
    "| model_name      | Pre-trained dataset         | in_channels |\n",
    "|-----------------|-----------------------------|-------------|\n",
    "| ViT-B-32        | laion400m_e31,laion400m_e32 | 512         |\n",
    "| ViT-B-16        | laion400m_e31               | 512         |\n",
    "| ViT-L-14        | laion400m_e31               | 768         |\n",
    "| ViT-H-14        | laion2b_s32b_b79k           | 1024        |\n",
    "| ViT-g-14        | laion2b_s12b_b42k           | 1024        |\n",
    "| EVA02-L-14      | merged2b_s4b_b131k          | 768         |\n",
    "| EVA02-L-14-336  | merged2b_s6b_b61k           | 768         |\n",
    "| EVA02-E-14      | laion2b_s4b_b115k           | 1024        |\n",
    "| EVA02-E-14-plus | laion2b_s9b_b144k           | 1024        |\n",
    "\n",
    "Please refer to license terms here [open_clip_license](https://huggingface.co/models?library=open_clip ) for all the open_clip models licensing.\n",
    "Please refer to the steps here to get a comprehensive list of models supported by OpenCLIP API: [open_clip models list](https://github.com/mlfoundations/open_clip#pretrained-models). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** If you are using the Logistic Regression head, the following parameters from the spec file model config should be used:\n",
    "\n",
    "* model:\n",
    "  * backbone:\n",
    "    * freeze: true\n",
    "    * pretrained: \"/path/to/NV_DINOV2_518.pth\"\n",
    "  * head:\n",
    "    * lr_head:\n",
    "      * C: 0.316   # tunable\n",
    "      * max_iter: 5000   # tunable\n",
    "    * type: LogisticRegressionHead\n",
    "    * num_classes: 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model -h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: The batch size has been set to 128 in the provided spec file. This needs a minimum of 20GB system memory. You can scale the BS based on your system mem. Recommended Batch size is 512 for best performance. Following is the Batch size vs System memory scaling:\n",
    "* 128 BS -> 20GB system mem\n",
    "* 256 BS -> 40GB system mem \n",
    "* 512 BS -> 80GB system mem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# This is the suitable number of epochs for this model with pretrained weights. Please change this value as needed.\n",
    "%env EPOCHS = 10\n",
    "%env NUM_GPUS = 8\n",
    "\n",
    "print(\"Train Classification Model\")\n",
    "!tao model classification_pyt train \\\n",
    "                  -e $SPECS_DIR/train_imagenet_clip.yaml \\\n",
    "                  results_dir=$RESULTS_DIR/classification_experiment_fm \\\n",
    "                  train.num_gpus=$NUM_GPUS \\\n",
    "                  train.train_config.runner.max_epochs=$EPOCHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"To resume from a checkpoint, use the below command. Update the epoch number accordingly\")\n",
    "# !tao model classification_pyt train \\\n",
    "#                   -e $SPECS_DIR/train_imagenet_clip.yaml \\\n",
    "#                   results_dir=$RESULTS_DIR/classification_experiment \\\n",
    "#                   train.num_gpus=$NUM_GPUS \\\n",
    "#                   train.train_config.resume_training_checkpoint_path=$RESULTS_DIR/classification_experiment/train/epoch_20.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('PyTorch checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/classification_experiment_fm/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "%env NUM_EPOCH=10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Rename a model: Note that the training is not deterministic, so you may change the model name accordingly.')\n",
    "print('---------------------')\n",
    "# NOTE: The following command may require `sudo`. You can run the command outside the notebook.\n",
    "!find $HOST_RESULTS_DIR/classification_experiment_fm/train -name \"epoch_$NUM_EPOCH.pth\" | xargs realpath | xargs -I {} mv {} $HOST_RESULTS_DIR/classification_experiment_fm/train/classification_model.pth \n",
    "!ls -ltrh $HOST_RESULTS_DIR/classification_experiment_fm/train/classification_model.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-4\"></a>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate ImageNet-1k Fine-tuned Classification Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_pyt evaluate \\\n",
    "                    -e $SPECS_DIR/test_clip_imagenet.yaml \\\n",
    "                    evaluate.num_gpus=$NUM_GPUS \\\n",
    "                    evaluate.checkpoint=$RESULTS_DIR/classification_experiment_fm/train/classification_model.pth \\\n",
    "                    results_dir=$RESULTS_DIR/classification_experiment_fm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inferences <a class=\"anchor\" id=\"head-5\"></a>\n",
    "In this section, we run the classification inference tool to generate inferences with the trained classification models and print the results. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model classification_pyt inference \\\n",
    "                    -e $SPECS_DIR/test_clip_imagenet.yaml \\\n",
    "                    inference.num_gpus=$NUM_GPUS \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/classification_experiment_fm/train/classification_model.pth \\\n",
    "                    results_dir=$RESULTS_DIR/classification_experiment_fm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the results\n",
    "!cat $HOST_RESULTS_DIR/classification_experiment_fm/inference/result.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the inference with images from the csv file. It contains the following columns - Image Name, class_label, class_confidence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install Deps\n",
    "!pip3 install pillow\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image \n",
    "import os\n",
    "import csv\n",
    "from math import ceil\n",
    "import random \n",
    "\n",
    "DATA_DIR = os.environ.get('HOST_DATA_DIR')\n",
    "DATA_DOWNLOAD_DIR = os.environ.get('DATA_DIR')\n",
    "RESULT_DIR = os.environ.get('HOST_RESULTS_DIR')\n",
    "csv_path = os.path.join(RESULT_DIR, \"classification_experiment_fm/inference/\" 'result.csv')\n",
    "results = []\n",
    "with open(csv_path) as csv_file:\n",
    "    csv_reader = csv.reader(csv_file, delimiter=',')\n",
    "    for row in csv_reader:\n",
    "        results.append((row[0], row[1]))\n",
    "random.shuffle(results)\n",
    "\n",
    "w,h = 200,200\n",
    "fig = plt.figure(figsize=(30,30))\n",
    "columns = 3\n",
    "rows = 1\n",
    "for i in range(1, columns*rows + 1):\n",
    "    ax = fig.add_subplot(rows, columns,i)\n",
    "    img = Image.open(results[i][0].replace(DATA_DOWNLOAD_DIR, DATA_DIR))\n",
    "    img = img.resize((w,h), Image.ANTIALIAS)\n",
    "    plt.imshow(img)\n",
    "    ax.set_title(results[i][1], fontsize=40)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Deploy! <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export the Classification model to ONNX model\n",
    "# NOTE: Export is done on single GPU - GPU num need not be provided\n",
    "\n",
    "!tao model classification_pyt export \\\n",
    "                   -e $SPECS_DIR/export_imagenet_clip.yaml \\\n",
    "                   export.checkpoint=$RESULTS_DIR/classification_experiment_fm/train/classification_model.pth \\\n",
    "                   export.onnx_file=$RESULTS_DIR/classification_experiment_fm/export/classification_model_export.onnx \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment_fm/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a TensorRT Engine using TAO Deploy\n",
    "!tao deploy classification_pyt gen_trt_engine \\\n",
    "                   -e $SPECS_DIR/export_imagenet_clip.yaml \\\n",
    "                   gen_trt_engine.onnx_file=$RESULTS_DIR/classification_experiment_fm/export/classification_model_export.onnx \\\n",
    "                   gen_trt_engine.trt_engine=$RESULTS_DIR/classification_experiment_fm/gen_trt_engine/classification_model_export.engine \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment_fm/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run evaluation using the generated TensorRT Engine\n",
    "!tao deploy classification_pyt evaluate \\\n",
    "                   -e $SPECS_DIR/test_clip_imagenet.yaml \\\n",
    "                   evaluate.trt_engine=$RESULTS_DIR/classification_experiment_fm/gen_trt_engine/classification_model_export.engine \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment_fm/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inference using the generated TensorRT Engine\n",
    "!tao deploy classification_pyt inference \\\n",
    "                   -e $SPECS_DIR/test_clip_imagenet.yaml \\\n",
    "                   inference.trt_engine=$RESULTS_DIR/classification_experiment_fm/gen_trt_engine/classification_model_export.engine \\\n",
    "                   results_dir=$RESULTS_DIR/classification_experiment_fm/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the results\n",
    "!cat $HOST_RESULTS_DIR/classification_experiment_fm/trt_inference/result.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize Inference\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image \n",
    "import os\n",
    "import csv\n",
    "from math import ceil\n",
    "import random \n",
    "\n",
    "DATA_DIR = os.environ.get('HOST_DATA_DIR')\n",
    "DATA_DOWNLOAD_DIR = os.environ.get('DATA_DIR')\n",
    "RESULT_DIR = os.environ.get('HOST_RESULTS_DIR')\n",
    "csv_path = os.path.join(RESULT_DIR, \"classification_experiment_fm/trt_inference/\" 'result.csv')\n",
    "results = []\n",
    "with open(csv_path) as csv_file:\n",
    "    csv_reader = csv.reader(csv_file, delimiter=',')\n",
    "    for row in csv_reader:\n",
    "        results.append((row[0], row[1]))\n",
    "random.shuffle(results)\n",
    "\n",
    "w,h = 200,200\n",
    "fig = plt.figure(figsize=(30,30))\n",
    "columns = 5\n",
    "rows = 1\n",
    "for i in range(1, columns*rows + 1):\n",
    "    ax = fig.add_subplot(rows, columns,i)\n",
    "    img = Image.open(results[i][0].replace(DATA_DOWNLOAD_DIR, DATA_DIR))\n",
    "    img = img.resize((w,h), Image.ANTIALIAS)\n",
    "    plt.imshow(img)\n",
    "    ax.set_title(results[i][1], fontsize=40)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:videollama] *",
   "language": "python",
   "name": "conda-env-videollama-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "8dd506d74b0e800f607ce4c2aea2a7d91bf4a69ac59ceef67481b8a20350139b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
